import { WorkflowType } from '@/components/workflow/types'
import { generateUUID } from '@/utils'
import { produce } from 'immer'
import { cloneDeep } from 'lodash'
import type { Node as ReactFlowNode } from 'reactflow'
import { ReactFlowInstance } from 'reactflow'
import { create } from 'zustand'
/**
 * 工作流存储
 */
type FlowVariablesType = {
  [key: string]: any
}

type WorkflowState = {
  flow: WorkflowType
  lastFlow: WorkflowType // 上一个自动更新前的flow信息，用来做自动保存时的数据比较，每次自动更新后更新该值
  updateTime: string | null
  runTime: string | null
  publishTime: string | null
  // reactFlow绑定的实例
  flowInstance: ReactFlowInstance | null
  selectedNode: any | null
  // 变量，输入输出在前端单独存储，回显和保存时根据id挂载到对应节点
  flowVariables: FlowVariablesType
  flowInputs: FlowVariablesType
  flowOutputs: FlowVariablesType
  // 校验
  verifying: boolean
  // 运行状态管理
  runType: string | 'single' | 'all'
  running: boolean
  runID: string
  singleRunningInfo: object | any
  errorInfo: string
  // 一些节点中需要使用的后端返回的可选项
  modelOptions: any[]
  // 禁止发布
  disabledPublish: boolean
  tool_list: any[]
  knowledge_list: any[]
  // 禁止编辑
  disabledEdit: boolean
}

type RunInfoType = {
  status?: string
  inputs?: string | any[]
  system_prompt?: string
  outputs?: string | any[]
  startTime?: number
  endTime?: number
  outputsList?: string | any[]
}

type WorkflowActions = {
  initFlow: (flow, tool_list?, knowledge_list?) => void
  destroyFlow: () => void
  updateLastFlow: (lastFlow) => void
  getFlowInstance: () => ReactFlowInstance | null
  setUpdateTime: (updateTime: string) => void
  setRunTime: (runTime: string) => void
  setPublishTime: (publishTime: string) => void
  setFlowInstance: (flowInstance: ReactFlowInstance) => void
  deleteNode: (idx: string) => void
  setSelectedNode: (node) => void
  setFlowVariables: (id: string, name: string, variables: object) => void
  coverFlowVariables: (variables: object) => void // 覆盖所有变量
  delFlowVariablesKey: (id: string) => void
  updateNodeAnimated: (runInfo) => void // 节点动画控制
  setVerifying: (status: boolean) => void
  updateRunType: (runType: 'single' | 'all') => void
  setRunning: (status: boolean) => void
  setRunID: (status: string) => void
  updateSingleNodeData: (nodeId, data) => void // 单个节点数据更新
  // 连接限制不采用dify设计逻辑，直接通过限制入口来控制连接，不限制出口，除结束节点外每个节点最多接收一个输入
  addNodeInIteration: (edge) => void // 节点连接建立
  removeNodesInIteration: (edges) => void // 节点连接移除
  removeNodesInIterationAllEdgs: (edges) => void // 节点 node连接移除
  updateSingleRunningInfo: (nodeId, running, info?: RunInfoType) => void // 单个节点运行状态
  updateRunningToError: () => void // 把运行中的节点状态修改为失败
  clearRunInfo: () => void // 清除所有运行状态及结果
  setErrorInfo: (text: string) => void // 设置错误信息
  setModelOptions: (options) => void // 设置大模型可选项
  setDisabledPublish: (status: boolean) => void
  setDisabledEdit: (status: boolean) => void
}

export const getInitFlowParams = () => {
  return {
    flow: null,
    lastFlow: null,
    updateTime: null,
    runTime: null,
    publishTime: null,
    flowInstance: null,
    selectedNode: null,
    flowVariables: {},
    flowInputs: {},
    flowOutputs: {},
    verifying: false,
    runType: 'all',
    running: false,
    runID: '',
    singleRunningInfo: {},
    errorInfo: '',
    modelOptions: [],
    disabledPublish: false,
    tool_list: [],
    knowledge_list: [],
    disabledEdit: false
  }
}

export const useWorkFlowStore = create<WorkflowState & WorkflowActions>((set, get) => ({
  ...getInitFlowParams(),
  initFlow: (flow, tool_list, knowledge_list) => {
    console.log('>>> initFlow: ', flow)
    set(state => ({ flow: { ...state.flow, ...flow } }))
    set(state => ({ lastFlow: cloneDeep({ ...state.flow, ...flow }) }))

    if (flow.update_time) {
      set({ updateTime: flow.update_time.replace('T', ' ') })
    }

    if (flow.publish_time) {
      set({ publishTime: flow.publish_time.replace('T', ' ') })
    }

    if (flow.run_timestamp) {
      set({ runTime: flow.run_timestamp.replace('T', ' ') })
    }
    if (tool_list) {
      set({ tool_list })
    }
    if (knowledge_list) {
      set({ knowledge_list })
    }
    // 存在节点时在初始化阶段提取所有节点的变量
    if (flow?.graph?.nodes?.length > 0) {
      const tempVariables = {}
      flow?.graph?.nodes?.forEach(item => {
        tempVariables[item.id] = item.data
      })

      console.log('>>> init variables tempVariables: ', tempVariables)
      set(state => ({ flowVariables: { ...state.flowVariables, ...tempVariables } }))
    }
  },
  destroyFlow() {
    console.log('>>> destroyFlow')
    set({ ...getInitFlowParams() })
  },

  updateLastFlow: lastFlow => {
    console.log('>>> updateLastFlow: ', lastFlow)
    set(state => ({ lastFlow: cloneDeep({ ...state.flow, ...lastFlow }) }))
  },
  getFlowInstance() {
    const flowInstance = get().flowInstance
    if (!flowInstance || Object.keys(flowInstance).length === 0) return null
    return flowInstance
  },
  setUpdateTime: (updateTime: string) => {
    set({ updateTime })
  },
  setRunTime: (runTime: string) => {
    set({ runTime })
  },
  setPublishTime: (publishTime: string) => {
    set({ publishTime })
  },
  setFlowInstance(flowInstance: ReactFlowInstance) {
    set({ flowInstance: { ...flowInstance } })
  },
  setSelectedNode: node => {
    console.log('>>> setSelectedNode node: ', node)
    if (node) {
      set({ selectedNode: { ...node } })
    } else {
      set({ selectedNode: null })
    }
  },
  deleteNode(id: string) {
    console.log('>>> deleteNode id: ', id)
    const flowInstance = get().flowInstance
    if (!flowInstance || Object.keys(flowInstance).length === 0) return

    // 获取所有节点
    const allNodes = flowInstance.getNodes()
    const nodeToDelete = allNodes.find(node => node.id === id)
    
    // 如果是迭代节点，找出所有子节点
    const nodesToDelete = [id]
    if (nodeToDelete?.data?.type === 'iteration' || nodeToDelete?.data?.type === 'loop') {
      allNodes.forEach(node => {
        if (node.parentId === id) {
          nodesToDelete.push(node.id)
        }
      })
    }

    // 获取要删除的边
    const edgesToDelete = flowInstance.getEdges().filter(edge => 
      nodesToDelete.includes(edge.source) || nodesToDelete.includes(edge.target)
    )

    // 重置受影响的迭代结束节点的状态
    const updatedNodes = allNodes.map(node => {
      if (node.data.type === 'iteration_end' || node.data.type === 'loop_end') {
        const isTargetOfDeletedEdge = edgesToDelete.some(edge => edge.target === node.id)
        const hasOtherIncomingEdges = flowInstance.getEdges().some(edge => 
          edge.target === node.id && !nodesToDelete.includes(edge.source)
        )
        
        if (isTargetOfDeletedEdge && !hasOtherIncomingEdges) {
          return {
            ...node,
            data: {
              ...node.data,
              isInIteration: false
            }
          }
        }
      }

      // 处理其他受影响的节点
      if (edgesToDelete.some(edge => edge.target === node.id || edge.source === node.id)) {
        return {
          ...node,
          data: {
            ...node.data,
            isInIteration: false,
          }
        }
      }

      return node
    })
    

    // 删除相关的边和节点
    flowInstance.setEdges(flowInstance.getEdges().filter(edge => 
      !nodesToDelete.includes(edge.source) && !nodesToDelete.includes(edge.target)
    ))
    flowInstance.setNodes(updatedNodes.filter(node => !nodesToDelete.includes(node.id)))

    // 移除相关节点的变量存储
    set(state => {
      const variables = get().flowVariables
      let validFlowVariables = variables
      nodesToDelete.forEach(nodeId => {
        validFlowVariables = removeFlowVariablesKey(nodeId, validFlowVariables)
      })

      // 清理引用了被删除节点的变量
      Object.entries(validFlowVariables).forEach(([key, value]) => {
        const hasDeletedNodeRef = value?.inputList?.some(item => 
          item?.value?.includes?.(id) || 
          item?.valuet?.includes?.(id)
        ) || value?.outputList?.some(item => 
          item?.value?.includes?.(id) || 
          item?.valuet?.includes?.(id)
        );

        if (hasDeletedNodeRef) {
          validFlowVariables[key] = {
            ...validFlowVariables[key],
            inputList: value?.inputList?.map(item => 
              item?.value?.includes?.(id) || item?.valueT?.includes?.(id)
                ? { ...item, value: '', valuet: '' }
                : item
            ) || [],
            outputList: value?.outputList?.map(item => 
              item?.value?.includes?.(id) || item?.valueT?.includes?.(id)
                ? { ...item, value: '', valuet: '' }
                : item
            ) || [],
          }
        }
      })

      return { flowVariables: validFlowVariables }
    })
  },
  // 参数配置信息更新
  setFlowVariables: (id: string, name: string, variables: object) => {
    console.log('>>> 变量缓存更新 id, variable: ', id, name, variables)
    set(state => {
      const newFlowVariables = produce(state.flowVariables, (draft: any) => {
        draft[id] = { ...draft[id], ...variables }
      })

      const flowInstance = get().flowInstance
      const validFlowVariables = removeFlowVariablesKeys(id, newFlowVariables, flowInstance)
      console.log('newFlowVariables, validFlowVariables', newFlowVariables, validFlowVariables)
      return { flowVariables: validFlowVariables }
    })
  },
  coverFlowVariables: (variables: object) => {
    console.log('>>> 覆盖变量信息 id, variable: ', variables)
    set(state => {
      return { flowVariables: variables }
    })
  },
  delFlowVariablesKey: (id: string) => {
    console.log('>>> 变量缓存删除 id : ', id)
    set(state => {
      const variables = get().flowVariables
      const validFlowVariables = removeFlowVariablesKey(id, variables)
      console.log('deleteVariables flowVariables, validFlowVariables', variables, validFlowVariables)
      return { flowVariables: validFlowVariables }
    })
  },
  setRunning: (status: boolean) => {
    set({ running: status })
  },

  updateNodeAnimated: runInfo => {
    const flowInstance = get().getFlowInstance()
    if (!flowInstance) return

    flowInstance.setEdges(
      flowInstance.getEdges().map(edge => {
        const targetStatus = runInfo[edge.target]?.status
        const isRunning = targetStatus === 'running'

        // 返回新的 edge 对象，包含 animated 属性
        return {
          ...edge,
          animated: isRunning
        }
      })
    )
  },
  updateRunType: (runType: 'single' | 'all') => {
    set({ runType })
  },
  setRunID: id => {
    set({ runID: id })
  },
  updateSingleNodeData: (nodeId, data) => {
    console.log('>>> 节点data数据更新 nodeId, status: ', nodeId, data)
    if (!nodeId || !data) return
    const flowInstance = get().getFlowInstance()
    if (!flowInstance) return

    flowInstance.setNodes(
      flowInstance
        .getNodes()
        .map((node: ReactFlowNode) => (node.id === nodeId ? { ...node, data: { ...node.data, ...data } } : node))
    )
  },
  addNodeInIteration: edge => {
    console.log('>>> 节点连接建立 edge: ', edge)
    if (!edge) return
    const flowInstance = get().getFlowInstance()
    if (!flowInstance) return

    const linkNodes = [edge.target]
    flowInstance.setNodes(
      flowInstance.getNodes().map((node: ReactFlowNode) => {
        return linkNodes.includes(node.id) ? { ...node, data: { ...node.data, isInIteration: true } } : node
      })
    )
  },
  removeNodesInIteration: edges => {
    if (!edges || edges.length === 0 || edges[0].type !== 'remove') return
    console.log('>>> 节点连接移除 edges: ', edges)

    const flowInstance = get().getFlowInstance()
    if (!flowInstance) return

    const removedNodes = edges.flatMap(edge => edge.id)
    const linkEdges = flowInstance.getEdges().filter(edge => !removedNodes.includes(edge.id))
    const linkNodes = linkEdges.flatMap(edge => [edge.target])
    flowInstance.setNodes(
      flowInstance.getNodes().map((node: ReactFlowNode) => {
        return linkNodes.includes(node.id)
          ? { ...node, data: { ...node.data, isInIteration: true } }
          : { ...node, data: { ...node.data, isInIteration: false } }
      })
    )
  },
  removeNodesInIterationAllEdgs: edges => {
    if (!edges || edges.length === 0 || edges[0].type !== 'remove') return
    console.log('>>> 节点连接移除 edges: ', edges)

    const flowInstance = get().getFlowInstance()
    if (!flowInstance) return

    // 获取需要移除的边ID列表
    const removedEdgeIds = edges.map(edge => edge.id)
    console.log('>>> 需要移除的边ID: ', removedEdgeIds)

    // 过滤出未被移除的边
    const remainingEdges = flowInstance.getEdges().filter(edge => !removedEdgeIds.includes(edge.id))
    console.log('>>> 剩余的边: ', remainingEdges)

    // 获取所有与剩余边相连的节点ID，避免重复，使用Set
    const removedNodes = edges.flatMap(edge => edge.id)

    const linkEdges = flowInstance.getEdges().filter(edge => !removedNodes.includes(edge.id))

    const linkNodes = linkEdges.flatMap(edge => [edge.target])

    flowInstance.setNodes(
      flowInstance.getNodes().map((node: ReactFlowNode) => {
        return linkNodes.includes(node.id)
          ? { ...node, data: { ...node.data, isInIteration: true } }
          : { ...node, data: { ...node.data, isInIteration: false } }
      })
    )

    // 使用setEdges更新边的集合，只保留未被移除的边
    flowInstance.setEdges(remainingEdges)
  },
  updateSingleRunningInfo: (nodeId, status, runInfo: RunInfoType = {}) => {
    console.log('>>> 节点运行状态更新 nodeId, status: ', nodeId, status, runInfo)
    const newSingleRunningInfo = produce(get().singleRunningInfo, draft => {
      const lastInfo = draft[nodeId] ? draft[nodeId] : {}
      draft[nodeId] = {
        ...lastInfo,
        status,
        ...runInfo
      }
    })
    set(state => {
      return { singleRunningInfo: newSingleRunningInfo }
    })
    // 更新连线动画
    get().updateNodeAnimated(newSingleRunningInfo)
  },
  updateRunningToError: () => {
    const newSingleRunningInfo = produce(get().singleRunningInfo, draft => {
      if (draft.status === 'running') {
        draft.status = 'failed'
      }
    })
    set({ singleRunningInfo: newSingleRunningInfo })
    // 更新连线动画
    get().updateNodeAnimated(newSingleRunningInfo)
  },
  clearRunInfo: () => {
    set(state => {
      return {
        verifying: false,
        running: false,
        singleRunningInfo: {}
      }
    })
  },
  setErrorInfo: text => {
    set({ errorInfo: text })
  },
  setVerifying: (status: boolean) => {
    set({ verifying: status })
  },
  setModelOptions: options => {
    set({ modelOptions: options })
  },
  setDisabledPublish: status => {
    set({ disabledPublish: status })
  },
  setDisabledEdit: status => {
    console.log('>>> setDisabledEdit status: ', status)
    set({ disabledEdit: status })
  }
}))

export function generateUniqueNodeId(nodeType: string) {
  return nodeType + '_' + generateUUID(5)
}

// 移除已删除节点的变量存储
export function removeFlowVariablesKeys(id, flowVariables, flowInstance) {
  const validFlowVariables = {}

  const ids = flowInstance
    .getNodes()
    .map(item => item.id)
    .concat([id])

  Object.keys(flowVariables).forEach(key => {
    if (ids.includes(key)) {
      validFlowVariables[key] = flowVariables[key]
    }
  })
  return validFlowVariables
}

// 移除指定ID的变量
export function removeFlowVariablesKey(id: string, flowVariables) {
  // 创建一个深拷贝，以确保我们不会修改原始对象
  const validFlowVariables = Object.fromEntries(Object.entries(cloneDeep(flowVariables)).filter(([key]) => key !== id))
  return validFlowVariables
}
